from datasets.mnist import MNIST
from utils.instant_tools import *
from utils.classifier import Classifier
from torch.utils.data import DataLoader
import torch.nn as nn
from models.fengyong_net import FengyongNet

BATCH_SIZE = 32
NUM_CATEGORIES = 10

if __name__ == "__main__":
    print("Loading MNIST dataset...")
    train_set = MNIST(is_train=True)
    test_set = MNIST(is_train=False)

    # 定义数据集加载器
    train_loader = DataLoader(train_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=10)
    test_loader = DataLoader(test_set, batch_size=BATCH_SIZE, shuffle=False, num_workers=10)

    # 定义网络模型
    backbone = FengyongNet(num_categories=NUM_CATEGORIES)
    model = torch.nn.Sequential(backbone, torch.nn.Softmax(dim=1))
    get_model_complexity(model, (1, 1, 28, 28))
    # 定义分类器，指定分类数和指标存储路径
    classifier = Classifier(model, num_category=NUM_CATEGORIES, criterion=nn.CrossEntropyLoss(),
                            metrics_save_folder="metrics", lr=0.05)
    # 训练和测试
    print("Training...")
    for epoch in range(30):
        print("Epoch: ", epoch)
        classifier.train(train_loader)
        if epoch > 0 and epoch % 5 == 0:
            classifier.test(test_loader)
